import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import akshare as ak
import time
from rich.console import Console
from rich.table import Table
from datetime import datetime, timedelta
import warnings
warnings.filterwarnings('ignore')
console = Console()
class SectorFlowMonitor:
    def __init__(self):
        """初始化板块资金流动监测器"""
        self.sector_data = {}
        self.sector_rank = {}
        self.historical_flow = {}
        
    def get_sector_list(self):
        """获取行业板块列表"""
        try:
            # 使用akshare获取行业板块数据
            industry_sector = ak.stock_board_industry_name_ths()
            
            # Convert to rich table and print
            # from rich.console import Console
            # from rich.table import Table
            # console = Console()
            table = Table(title="Industry Sector")
            
            # Add columns
            for col in industry_sector.columns:
                table.add_column(col)
            
            # Add rows
            for index, row in industry_sector.iterrows():
                table.add_row(*[str(item) for item in row])
            
            console.print(table)
            return industry_sector['code'].tolist()
        except Exception as e:
            print(f"获取板块列表失败: {e}")
            return []
    
    def get_sector_fund_flow(self, sector_name=None, date=None):
        """
        获取板块资金流向数据
        :param sector_name: 板块名称，None表示获取所有板块
        :param date: 查询日期，格式：YYYY-MM-DD，None表示当前日期
        :return: 板块资金流向数据
        """
        if date is None:
            date = datetime.now().strftime('%Y-%m-%d')
        
        try:
            # 使用akshare获取板块资金流向
            if sector_name:
                # 获取单个板块资金流向
                flow_data = ak.stock_individual_fund_flow(stock=sector_name, market="sh")
                return flow_data
            # else:
                # 获取所有板块资金流向
                # all_sector_flow = ak.stock_sector_fund_flow(indicator="行业资金流", market="sh")
                # return all_sector_flow
        except Exception as e:
            print(f"获取{sector_name}板块资金流向失败: {e}")
            return pd.DataFrame()
    
    def monitor_sector_flow(self, top_n=10):
        """
        监测当前资金流入流出最大的板块
        :param top_n: 显示前n个板块
        """
        # 获取所有板块资金流向
        all_flow = self.get_sector_fund_flow()
        
        if all_flow.empty:
            print("未能获取板块资金流向数据")
            return
        
        # 保存当前数据
        current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        self.sector_data[current_time] = all_flow
        
        # 按照主力净流入排序
        all_flow['主力净流入-净额'] = pd.to_numeric(all_flow['主力净流入-净额'], errors='coerce')
        inflow_sectors = all_flow.sort_values(by='主力净流入-净额', ascending=False).head(top_n)
        outflow_sectors = all_flow.sort_values(by='主力净流入-净额', ascending=True).head(top_n)
        
        # 保存排名
        self.sector_rank[current_time] = {
            'inflow': inflow_sectors,
            'outflow': outflow_sectors
        }
        
        print(f"\n{current_time} 资金流向监测:")
        print("主力资金净流入最大的板块:")
        for i, (_, row) in enumerate(inflow_sectors.iterrows(), 1):
            print(f"{i}. {row['板块名称']}: {row['主力净流入-净额']}万元")
        
        print("\n主力资金净流出最大的板块:")
        for i, (_, row) in enumerate(outflow_sectors.iterrows(), 1):
            print(f"{i}. {row['板块名称']}: {row['主力净流入-净额']}万元")
        
        return inflow_sectors, outflow_sectors
    
    def track_sector_movement(self, sector_names, days=5):
        """
        跟踪特定板块的资金流动趋势
        :param sector_names: 板块名称列表
        :param days: 回溯天数
        """
        for sector in sector_names:
            sector_flow_data = []
            start_date = (datetime.now() - timedelta(days=days)).strftime('%Y%m%d')
            
            # 这里简化处理，实际应用中需要按天获取数据
            print(f"\n获取{sector}板块{days}天资金流动数据...")
            
            # 为演示，使用当前数据模拟多天数据
            for i in range(days):
                date = (datetime.now() - timedelta(days=i)).strftime('%Y-%m-%d')
                daily_flow = self.get_sector_fund_flow(sector)
                
                if not daily_flow.empty:
                    # 提取需要的数据
                    daily_flow = daily_flow.iloc[0] if isinstance(daily_flow, pd.DataFrame) else daily_flow
                    flow_record = {
                        'date': date,
                        'sector': sector,
                        'net_inflow': float(daily_flow.get('主力净流入-净额', 0)),
                        'large_inflow': float(daily_flow.get('超大单净流入-净额', 0)),
                        'medium_inflow': float(daily_flow.get('大单净流入-净额', 0)),
                        'small_inflow': float(daily_flow.get('中单净流入-净额', 0))
                    }
                    sector_flow_data.append(flow_record)
            
            if sector_flow_data:
                # 保存历史数据
                self.historical_flow[sector] = pd.DataFrame(sector_flow_data)
                self.visualize_sector_flow(sector, self.historical_flow[sector])
    
    def visualize_sector_flow(self, sector_name, flow_data):
        """
        可视化特定板块的资金流动
        :param sector_name: 板块名称
        :param flow_data: 资金流动数据
        """
        if flow_data.empty:
            print(f"没有{sector_name}的资金流动数据可可视化")
            return
        
        # 按日期排序
        flow_data = flow_data.sort_values('date')
        
        plt.figure(figsize=(12, 6))
        
        # 绘制主力资金净流入
        plt.subplot(2, 1, 1)
        bars = plt.bar(flow_data['date'], flow_data['net_inflow'], color='blue')
        plt.title(f'{sector_name} 主力资金净流入趋势')
        plt.ylabel('净流入(万元)')
        
        # 添加数值标签
        for bar in bars:
            height = bar.get_height()
            plt.text(bar.get_x() + bar.get_width()/2., height,
                     f'{height:.2f}',
                     ha='center', va='bottom', rotation=0)
        
        # 绘制资金分类流入
        plt.subplot(2, 1, 2)
        x = range(len(flow_data['date']))
        plt.plot(x, flow_data['large_inflow'], 'r-', label='超大单')
        plt.plot(x, flow_data['medium_inflow'], 'g-', label='大单')
        plt.plot(x, flow_data['small_inflow'], 'b-', label='中单')
        plt.xticks(x, flow_data['date'])
        plt.title(f'{sector_name} 不同规模资金净流入')
        plt.ylabel('净流入(万元)')
        plt.legend()
        
        plt.tight_layout()
        plt.show()
    
    def analyze_sector_rotation(self, period=5):
        """
        分析板块轮动情况
        :param period: 分析周期(天)
        :return: 板块轮动分析结果
        """
        print(f"\n分析最近{period}天板块轮动情况...")
        
        # 获取最近period天的板块数据
        recent_dates = sorted(self.sector_data.keys())[-period:]
        if len(recent_dates) < 2:
            print("数据不足，无法分析板块轮动")
            return
        
        rotation_analysis = {}
        
        # 分析每天资金流向变化
        for i in range(1, len(recent_dates)):
            prev_date = recent_dates[i-1]
            curr_date = recent_dates[i]
            
            prev_data = self.sector_data[prev_date]
            curr_data = self.sector_data[curr_date]
            
            # 合并数据以便比较
            merged_data = pd.merge(
                prev_data[['板块名称', '主力净流入-净额']],
                curr_data[['板块名称', '主力净流入-净额']],
                on='板块名称',
                suffixes=('_prev', '_curr')
            )
            
            # 计算资金流向变化
            merged_data['flow_change'] = merged_data['主力净流入-净额_curr'] - merged_data['主力净流入-净额_prev']
            
            # 找出变化最大的板块
            top_increase = merged_data.sort_values('flow_change', ascending=False).head(5)
            top_decrease = merged_data.sort_values('flow_change', ascending=True).head(5)
            
            rotation_analysis[f'{prev_date}至{curr_date}'] = {
                'increase': top_increase,
                'decrease': top_decrease
            }
            
            print(f"\n{prev_date}至{curr_date}资金流向变化:")
            print("资金流入增加最大的板块:")
            for _, row in top_increase.iterrows():
                print(f"  {row['板块名称']}: +{row['flow_change']:.2f}万元")
            
            print("资金流入减少最大的板块:")
            for _, row in top_decrease.iterrows():
                print(f"  {row['板块名称']}: {row['flow_change']:.2f}万元")
        
        return rotation_analysis


def main():
    # 创建板块资金流动监测器
    monitor = SectorFlowMonitor()
    
    # 获取板块列表
    sectors = monitor.get_sector_list()
    if not sectors:
        print("无法获取板块列表，程序退出")
        return
    
    print(f"成功获取{len(sectors)}个行业板块")
    
    # 1. 实时监测资金流向
    print("\n===== 实时资金流向监测 =====")
    # top_inflow, top_outflow = monitor.monitor_sector_flow(top_n=5)
    
    # 2. 跟踪特定板块资金流动
    print("\n===== 板块资金流动趋势跟踪 =====")
    # tracked_sectors = ['银行', '房地产', '半导体', '新能源', '医药']

    monitor.track_sector_movement(sectors, days=5)
    
    # 3. 分析板块轮动
    print("\n===== 板块轮动分析 =====")
    # 为演示，模拟多天数据
    for i in range(5):
        monitor.monitor_sector_flow(top_n=5)
        time.sleep(1)  # 模拟不同时间点
    
    monitor.analyze_sector_rotation(period=3)


if __name__ == "__main__":
    main()    
